import numpy as np
import scipy.interpolate as interp
from scipy.signal import savgol_filter as sf


pi=np.pi
mrange=100
size=64
n1=40
n=10
m=0

m1=0
m2=0
m3=0
# cardino coordinate 
x=y=z=np.linspace(-50,50,size,dtype='float32')-0.7936508
x,y,z = np.meshgrid(x,y,z)
x=x.reshape(size**3,1)
y=y.reshape(size**3,1)
z=z.reshape(size**3,1)

# calculate the influence range at each x, y, z position
beta=np.arctan(z/(np.sqrt(x**2+y**2)+1e-5))
alpha=np.arctan(y/(x+1e-5))

# calculate the distance between given position to injection position       
radius=np.sqrt(x**2+y**2+z**2)

R=np.random.rand(n1,3)
R[R<0.01]=1.0

for k in range(2):
    if k==1: n=n1 
    for i in range(n):
        ac,bc,cc=R[i]
        a=ac*mrange
        b=bc*mrange
        c=cc*mrange
        coef=np.square(np.cos(alpha)*np.cos(beta)/a)+np.square(np.sin(alpha)*np.cos(beta)/b)+np.square(np.sin(beta)/c)
        xyz_range=np.sqrt(1/coef)         
 
        # randomize the influence range in beta and alpha angle delta=0.1 5o
        if k==1:   
            n_beta=np.random.randint(5,10,1)[0]
            n_alpha=np.random.randint(5,10,1)[0]
            x_beta=np.linspace(-pi/2,pi/2,n_beta)
            x_alpha=np.linspace(-pi/2,pi/2,n_alpha)        
            xx,yy=np.meshgrid(x_beta, x_alpha)
            Ratio=0.5-np.random.rand(x_alpha.size,x_beta.size)
            
            Ratio_s=sf(Ratio.flatten(),15,3).reshape(x_alpha.size,x_beta.size)
            
            f=interp.interp2d(x_beta,x_alpha,Ratio_s) 
            
            R_range=np.zeros([alpha.size,1])
            for idex in range (alpha.size):
                R_range[idex]=f(beta[idex],alpha[idex])      

            xyz_range=xyz_range+xyz_range* R_range.reshape(R_range.size,1)  
            

        xyz_range[xyz_range<10.]=10.
        xyz_range[xyz_range>50.]=50.

 
        q=np.random.randint(1,4,1)
        if q==1:
            v=np.exp(-5*(radius**2/xyz_range**2)) # divided by 5 to make v value close to zero at given xyz range  
            m1+=1
            print("Gaussian:",m1)
        elif q==2:
            v=np.exp(-5*(radius/xyz_range))    
            m2+=1
            print("exponential",m2)
        else:
            m3+=1
            v=1-radius/xyz_range # linear model
            print("linear",m3)
        
        # add backgroundwater noise of -0.1 to 0.1  
        randnumber=0.1-0.2*(-1+2*np.random.rand(size**3,1))
        v=v+randnumber*v
        v[v<1e-3]=0.0
        v[v>1]=1.0
        
        m=m+1
        np.savez_compressed('train_data_{00000}'.format(m), q=v)
        

